#!/usr/bin/env python
# encoding: utf-8

from source.Lib.Tools import Outer, StockInfo
from source.QuotaCalc import tool_function as tool
import datetime
from sqlalchemy.sql import func
from source.Lib import orm
import math
import multiprocessing
import json
import re


def get_list_max_count(string):
    if is_list_string(string):
        key_word, right_word = string.split('[')
        max_count = 0
        for i in right_word[:right_word.find("]")].split(':'):
            number = int(i) if i else 0
            if number < 0:
                temp = number * -1
            else:
                temp = number
            if temp > max_count:
                max_count = temp
    else:
        max_count = 1
    return max_count


def is_list_string(string):
    flag = len({'[', ']', ':'} - set(string))
    if flag == 0:
        return True
    else:
        return False


def calc_formula(code, date, formula, sources):
    outer = Outer()
    for field in sources:
        if "@" in field:
            formula = formula.replace(field, str(outer.get(f"{field}:{code}")))
        elif "$" in field:
            formula = formula.replace(field, str(QuotaUpdateSignal.GetReportQuota(field, code)))
        elif "&" in field:
            markets = list()
            max_count = get_list_max_count(field)
            for i in range(0, max_count*5):
                market = QuotaUpdateSignal.Get(
                    f"StockDay:{code}:{(date - datetime.timedelta(days=i)).strftime('%Y-%m-%d')}")
                if market:
                    markets.append(market)
                if len(markets) == max_count:
                    break

            if is_list_string(field):
                col_index = int(re.findall(r"&(\d+)", field)[0])
                target = f"[{','.join(list(map(lambda x: x.split(',')[col_index], markets)))}]"
                value = re.sub(r"(&.*?\])", target, field)
                market_pos = re.findall(r"&(\d+)", field)[0]
                source = field.replace("&" + market_pos, "&" + QuotaUpdateSignal.Get(f"MarketKeyWord:{market_pos}"))
                formula = formula.replace(source, value)
            else:
                formula = formula.replace(field, markets[0].split(',')[int(field.replace('&', ''))])
    return code, date, eval(formula) if {'&', '$', '@'} - set(formula.split()) else None


class Quota(object):
    """
        1. 根据公式自动生成计算函数:
            支持操作符: +, -, *, /, %, **, ()
            内置替换变量:
                &POS : 行情的字段的下标  [名称，开盘，收盘，当前，最高，最低，买一，卖一，成交股数，成交金额，买一手，买一价...卖一手，卖一价...]
                ${NAME}: 报表的指标的名称
                @ID  : 当前指标表的指标
            DEMO:
                =${每股收益} / ${开盘价}
                =${交易金额} * @123
                =${每股利润}
        2. 字典保存关注的指标, 加入设置接口
    """
    def __init__(self, qid, qname, qformula, sources, market):
        self.id = qid
        self.name = qname
        self.formula = qformula
        self.sources = sources
        # 注册指标更新信号，这个必须
        QuotaUpdateSignal.Register(self, sources)

        # 向行情信号注册自身，根据market类型
        if market:
            MarketUpdateSignal.Register(self)

    def calc(self, code, date):
        _, _, ret = calc_formula(code, date, self.formula, self.sources)
        return ret

    def flush(self, code, mode):
        # mode in {now(现在), all(所有), 2018-01-01(指定日期)}
        mem_data = QuotaUpdateSignal.Get(f"QuotaData:{self.id}:{code}")
        if not mem_data:
            mem_data = dict()
        else:
            mem_data = json.loads(mem_data.replace("'", '"'))

        if mode not in ('all', 'recal'):
            date = datetime.datetime.now() if mode == 'now' else datetime.datetime.strptime(mode, "%Y-%m-%d")
            if date.strftime('%Y-%m-%d') not in mem_data:
                ret = self.calc(code, date)
                if ret is not None:
                    QuotaUpdateSignal.Set(f"@{self.id}:{code}", ret)
                    QuotaUpdateSignal.Set(f"@{self.name}:{code}", ret)
                    QuotaUpdateSignal.Send(f"@{self.id}:{code}")
                    mem_data[date.strftime('%Y-%m-%d')] = round(ret, 1)
                    QuotaUpdateSignal.Set(f"QuotaData:{self.id}:{code}", json.dumps(mem_data))
        else:
            result = list()
            pool = multiprocessing.Pool(3)
            date = datetime.datetime.now()
            jump_day = 0
            if 'recal' == mode:
                mem_data = {}

            while date.strftime("%Y-%m-%d") not in mem_data and jump_day < 30:
                if not QuotaUpdateSignal.Get(f"StockDay:{code}:{date.strftime('%Y-%m-%d')}"):
                    jump_day += 1
                else:
                    jump_day = 0
                result.append(pool.apply_async(calc_formula, (code, date, self.formula, self.sources)))
                date -= datetime.timedelta(days=1)

            pool.close()
            pool.join()
            for ret in result:
                _, date, value = ret.get()
                if value:
                    mem_data[date.strftime('%Y-%m-%d')] = round(value, 1)
                    
            QuotaUpdateSignal.Set(f"QuotaData:{self.id}:{code}", json.dumps(mem_data))

            pool.close()
            pool.join()
            print(f"Quota:{self.id}, Code:{code}\r")


class MarketUpdateSignal(object):
    """
        - 接收消息队列的信号，读取其中Code。
        - 单例模型, 在Quota类中调用Register方式，注册指标， 指标=计算公式
            - 当信号进入，将触发计算公式, 计算完以后，由QuotaUpdateSignal发送指标更新信号
    """
    outer = Outer()
    register_list = set()

    @classmethod
    def Register(cls, quota):
        cls.register_list.add(quota)

    @classmethod
    def Start(cls):
        ps = cls.outer.get_ps()
        for item in ps.listen():
            if item['type'] == 'message':
                try:
                    quota_id, code, mode = bytes(item['data']).decode().split(':')
                except (Exception, ):
                    code, mode = bytes(item['data']).decode().split(':')
                    quota_id = ""

                for quota in cls.register_list:
                    if not quota_id or int(quota.id) == int(quota_id):
                        try:
                            if code != 'all':
                                quota.flush(code, mode)
                            else:
                                session = orm.Session()
                                codes = list(map(lambda x: x[0], session.query(orm.StockInfo.code).all()))
                                session.close()
                                for code in codes:
                                    quota.flush(code, mode)
                        except (TypeError, ):
                            pass


class QuotaUpdateSignal(object):
    """
        - 操作Redis缓存
        - 发送更新信号
    """
    outer = Outer(channel="QuotaUpdateSignal")
    stock = StockInfo()
    memory = {
        "report": dict(),
        "primary": dict()
    }
    dtRegister = dict()

    @classmethod
    def GetReportQuota(cls, k, code):
        """
            :param k:  $ + ID
            :param code: 股票代码
            :return: 返回数值 或者 None
        """
        data = json.loads(QuotaUpdateSignal.Get(f"Report:{code}").replace("'", '"'))
        quota_name = QuotaUpdateSignal.Get(f"ReportQuota:{k[1:]}").split(',')[-1]
        return data.get(quota_name, None)

    @classmethod
    def Send(cls, msg):
        cls.outer.publish(msg)

    @classmethod
    def Set(cls, k, v):
        cls.outer.put(k, v)

    @classmethod
    def Keys(cls, k):
        return cls.outer.keys(k)

    @classmethod
    def Get(cls, k):
        return cls.outer.get(k)

    @classmethod
    def Register(cls, qObject, sources):
        if sources:
            for s in sources:
                cls.dtRegister.setdefault(s, set())
                cls.dtRegister[s].add(qObject)

    @classmethod
    def Start(cls):
        ps = cls.outer.get_ps()
        for item in ps.listen():
            if item['type'] == 'message':
                name, code = bytes.decode(item['data']).split(':')
                for qObject in cls.dtRegister.get(name, []):
                    qObject.flush(code, "now")


if __name__ == '__main__':
    print(QuotaUpdateSignal.GetReportQuota("$1", "300160"))
